# from doc.pycurl.examples.quickstart.response_headers import headers
from flask import Flask, request, Response, jsonify, render_template, session
from neo4j import GraphDatabase
# from numba.cuda import stream
from werkzeug.security import generate_password_hash, check_password_hash
import pymysql.cursors
import os
import re
from functools import wraps
from flask import redirect, url_for
import pandas as pd
from collections import OrderedDict
import jieba
from fuzzywuzzy import fuzz
import json



# import ollama

# pip install flask neo4j pymysql jieba ollama fuzzywuzzy 

#调用chatglm
from zhipuai import ZhipuAI, APIStatusError

client = ZhipuAI(api_key="4318d8b2a8d1e74985d9066cf5251df5.YKEckaLOukDmI9A0") # 填写您自己的APIKey

conversation = []

app = Flask(__name__)
import json
import requests
# from rapidfuzz import process, fuzz
# from sentence_transformers import SentenceTransformer, util
# from sklearn.feature_extraction.text import TfidfVectorizer
# from sklearn.metrics.pairwise import cosine_similarity
# from sklearn.feature_extraction.text import CountVectorizer
# from sklearn.metrics import jaccard_score

with open('./static/config/config.json', encoding='utf-8') as conf:
    config = json.load(conf)

app.secret_key = os.getenv('SECRET_KEY', '123456')

uri = f"bolt://{config['host']}:7687"
driver = GraphDatabase.driver(uri, auth=(config['neo4j_user'], config['neo4j_password']))  # 使用你的用户名和密码

# 配置数据库连接
db_config = {
    'host': os.getenv('DB_HOST', config['mysql_url']),
    'port': int(os.getenv('DB_PORT', 3306)),
    'user': os.getenv('DB_USER', config['mysql_user']),
    'password': os.getenv('DB_PASSWORD', config['mysql_password']),
    'db': os.getenv('DB_NAME', config['database']),
    'charset': 'utf8mb4',
    'cursorclass': pymysql.cursors.DictCursor
}


# 创建数据库连接函数
class DBConnection:
    def __enter__(self):
        self.conn = pymysql.connect(**db_config)
        self.cursor = self.conn.cursor()
        return self.cursor

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.conn.commit()
        self.cursor.close()
        self.conn.close()


# 装饰器检查用户是否登录
def login_required(f):
    @wraps(f)
    def decorated_function(*args, **kwargs):
        if 'user_id' not in session:
            return redirect(url_for('index'))
        return f(*args, **kwargs)

    return decorated_function


@app.route('/check_login_status', methods=['GET'])
def check_login_status():
    if 'user_id' in session:
        return jsonify({"logged_in": True, "userRole": session.get('userRole', 'No role assigned')}), 200
    else:
        return jsonify({"logged_in": False}), 200


@app.route('/logout')
def logout():
    session.pop('user_id', None)
    session.pop('userRole', None)
    return jsonify({"success": "Logged out successfully"}), 200


@app.route('/login', methods=['POST'])
def login():
    username = request.form.get('username')
    password = request.form.get('password')
    userRole = request.form.get('userRole')
    if not username or not password:
        return jsonify({"error": "Username and password are required"}), 400

    try:
        with DBConnection() as cursor:
            cursor.execute("SELECT id, password, role FROM t_users WHERE username = %s", (username,))
            user = cursor.fetchone()
        if user and user["role"] == userRole:
            if user and check_password_hash(user['password'], password):
                session['user_id'] = user['id']
                session['userRole'] = user['role']
                return jsonify({"success": "Logged in successfully"}), 200

        return jsonify({"error": "Invalid username or password"}), 401
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@app.route('/register', methods=['POST'])
def register():
    username = request.form.get('username')
    password = request.form.get('password')
    userRole = request.form.get('userRole')
    if not username or not password:
        return jsonify({"error": "Username and password are required"}), 400

    try:
        with DBConnection() as cursor:
            cursor.execute("SELECT id FROM t_users WHERE username = %s", (username,))
            if cursor.fetchone():
                return jsonify({"error": "Username already exists"}), 409

            password_hash = generate_password_hash(password)
            cursor.execute("""
                INSERT INTO t_users (username, password, email, full_name, clinical_title,
                                   teaching_title, department, expertise, professional_experience, research_achievements, role)
                VALUES (%s, %s, '', '', '', '', '', '', '', '', %s)
            """, (username, password_hash, userRole))

        return jsonify({"success": "User registered successfully"}), 201
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@app.route('/change_password', methods=['POST'])
@login_required
def change_password():
    user_id = session.get('user_id')
    current_password = request.form.get('current_password')
    new_password = request.form.get('new_password')

    if not current_password or not new_password:
        return jsonify({"error": "Current password and new password are required"}), 400

    try:
        with DBConnection() as cursor:
            cursor.execute("SELECT password FROM t_users WHERE id = %s", (user_id,))
            user = cursor.fetchone()

            if not user or not check_password_hash(user['password'], current_password):
                return jsonify({"error": "Current password is incorrect"}), 401

            new_password_hash = generate_password_hash(new_password)
            cursor.execute("UPDATE t_users SET password = %s WHERE id = %s", (new_password_hash, user_id))

        return jsonify({"success": "Password changed successfully"}), 200
    except Exception as e:
        # log_exception(e)
        return jsonify({"error": str(e)}), 500


@app.route('/get_profile', methods=['GET'])
def get_profile():
    user_id = session.get('user_id')
    if not user_id:
        return jsonify({"error": "Not authenticated"}), 403

    try:
        with DBConnection() as cursor:
            cursor.execute("""
                SELECT email, full_name, clinical_title, teaching_title, department
                FROM t_users WHERE id = %s
            """, (user_id,))
            user = cursor.fetchone()

        if user:
            return jsonify({
                "success": True,
                "profile": {
                    "email": user['email'],
                    "full_name": user['full_name'],
                    "clinical_title": user['clinical_title'],
                    "teaching_title": user['teaching_title'],
                    "department": user['department']
                    # "expertise": user['expertise'],
                    # "professional_experience": user['professional_experience'],
                    # "research_achievements": user['research_achievements']
                }
            })
        else:
            return jsonify({"error": "User not found"}), 404
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@app.route('/update_profile', methods=['POST'])
def update_profile():
    user_id = session.get('user_id')
    if not user_id:
        return jsonify({"error": "Not authenticated"}), 403

    email = request.form.get('email', '')
    full_name = request.form.get('full_name', '')
    clinical_title = request.form.get('clinical_title', '')
    teaching_title = request.form.get('teaching_title', '')
    department = request.form.get('department', '')
    # expertise = request.form.get('expertise', '')
    # professional_experience = request.form.get('professional_experience', '')
    # research_achievements = request.form.get('research_achievements', '')

    try:
        with DBConnection() as cursor:
            cursor.execute("""
                UPDATE t_users SET
                email = %s, full_name = %s, clinical_title = %s, teaching_title = %s, 
                department = %s
                WHERE id = %s
            """, (email, full_name, clinical_title, teaching_title, department, user_id))

        return jsonify({"success": "Profile updated successfully"}), 200
    except Exception as e:
        return jsonify({"error": str(e)}), 500

def get_medical_words_through_models(ip, content, model):
    # if config["if_ip"] != "true":
    #     return 0
    if len(content) < 7:
        return content
    # url = f'http://{ip}:11434/api/chat'
    # url = f"{config['base_url']}/chat"
    url = config['base_url']
    full_content = f"请用中文回答，我现在会给你一句话，这句话会放在‘’之中, 请你从这句话中提取医学相关术语，请尽可能提取,提取内容要在原句中可查到，并将医学相关术语逐个给出，中间以顿号分隔，回答只有给出医学术语即可。要提取的话：‘{content}‘"
    data = {
        "model": model,
        "messages": [
            {
                "role": "user",
                "content": full_content
            }
        ]
    }
    headers = config['headers']
    try:
        response = requests.post(url=url, data=json.dumps(data), headers=headers)
        model_response = []
        for line in response.text.strip().splitlines():
            try:
                json_line = json.loads(line)  # 解析每一行的 JSON
                if 'message' in json_line and 'content' in json_line['message']:
                    model_response.append(json_line['message']['content'])
            except json.JSONDecodeError as json_error:
                print("JSON 解码错误:", json_error)
        final_response = ''.join(model_response)  # 或者使用 ' '.join(model_response) 以空格分隔
        return final_response if final_response else 0
    except requests.exceptions.RequestException as e:
        print("!!!!!!!!!!!!!!!RequestException!!!!!!!!!!!!!!!")
        return 0

# 获取图数据函数
def get_graph_data(name, limit=20):
    methods_of_div_sentence = "jieba"                
    nodes = []
    relationships = []
    medical_words = get_medical_words_through_models(config["ip"], name, config["model"])
    medical_words = medical_words if medical_words else name
    medical_words = medical_words.replace(' ','')
    medical_words_list = get_sub_words(medical_words) if methods_of_div_sentence == "jieba" else medical_words.split('、')
    
    for sub_search_name in medical_words_list:
        try:
            with driver.session() as session:
                query = """
                MATCH (a)-[r]-(b)
                WHERE a.name CONTAINS $name
                RETURN a, r, b LIMIT 128
                """
                result = session.run(query, name=sub_search_name)
                if not result:
                    continue
                for record in result:
                    a = record['a']
                    b = record['b']
                    r = record['r']
                    if a.id not in [node['id'] for node in nodes]:
                        nodes.append({
                            "id": a.id,
                            "types": a["types"],
                            "properties": {
                                "name": a._properties.get("name"),
                                "description": a._properties.get("description", "")  # 获取简介
                            }
                        })
                    if b.id not in [node['id'] for node in nodes]:
                        nodes.append({
                            "id": b.id,
                            "types": b["types"],
                            "properties": {
                                "name": b._properties.get("name"),
                                "description": b._properties.get("description", "")  # 获取简介
                            }
                        })
                    relationships.append({
                        "id": r.id,
                        "type": r.type,
                        "start": r.start_node.id,
                        "end": r.end_node.id,
                        "properties": r._properties
                    })
        except Exception as e:
            return {"error": str(e)}
    results = sorted([(node, fuzz.ratio(node["properties"]["name"], medical_words)) for node in nodes], key=lambda x: x[1], reverse=True)[:limit]
    nodes = [result[0] for result in results]
    return {"nodes": nodes, "relationships": relationships}


def expand_node(node_id):
    try:
        with driver.session() as session:
            result = session.run(
                "MATCH (a)-[r]-(b) WHERE id(a) = $id "
                "RETURN a, r, b LIMIT 128",
                id=node_id
            )
            nodes = []
            relationships = []
            for record in result:
                a = record['a']
                b = record['b']
                r = record['r']
                if a.id not in [node['id'] for node in nodes]:
                    nodes.append({"id": a.id, "types": a["types"], "properties": a._properties})
                if b.id not in [node['id'] for node in nodes]:
                    nodes.append({"id": b.id, "types": b["types"], "properties": b._properties})
                relationships.append({"id": r.id, "type": r.type, "start": r.start_node.id, "end": r.end_node.id,
                                      "properties": r._properties})

            return {"nodes": nodes, "relationships": relationships}
    except Exception as e:
        return {"error": str(e)}

@app.route('/default-graph', methods=['GET'])
def default_graph():
    # default_name = "妊娠心脏病"
    # data = get_graph_data(default_name,limit=1)
    default_name = "1级"
    data = get_graph_data(default_name)
    return jsonify(data)


@app.route('/')
def index():
    return render_template('index.html')

@app.route('/page')
@login_required
def page():
    return render_template('page.html')

@app.route('/knowledge')
@login_required
def knowledge():
    return render_template('knowledge.html')


@app.route('/knowledge_admin')
@login_required
def knowledge_admin():
    return render_template('knowledge_admin.html')


@app.route('/profile')
@login_required
def profile():
    return render_template('profile.html')


@app.route('/query', methods=['GET'])
def query():
    name = request.args.get('name')
    if not name:
        return jsonify({"error": "Name parameter is required"}), 400
    data = get_graph_data(name)
    return jsonify(data)


@app.route('/expand', methods=['GET'])
def expand():
    node_id = int(request.args.get('id'))
    data = expand_node(node_id)
    return jsonify(data)

def get_sub_words(sentence):
    return jieba.lcut(sentence, cut_all=False)

# def get_result_using_trans_model(search_name, string_list, limit):
#     model = SentenceTransformer('distiluse-base-multilingual-cased-v2')
#     encoded_search_name, encoded_string_list = model.encode(search_name, convert_to_tensor=True), model.encode(string_list, convert_to_tensor=True)
#     sorted_list_with_similarities = sorted(zip(string_list, util.pytorch_cos_sim(encoded_search_name, encoded_string_list)[0]), key=lambda x: x[1], reverse=True)
#     results = [item[0] for item in sorted_list_with_similarities]
#     return results[:limit]

# def get_most_similar_strings(search_name, string_list, limit):
#     vectorizer = CountVectorizer(analyzer='char', ngram_range=(2, 2))  # 使用2-gram字符模型
#     X = vectorizer.fit_transform(string_list)
#     search_vec = vectorizer.transform([search_name])
#     similarities = []
#     for i in range(len(string_list)):
#         similarity = jaccard_score(search_vec.toarray()[0], X.toarray()[i], average='binary')
#         similarities.append(similarity)
#     sorted_strings = [x for _, x in sorted(zip(similarities, string_list), reverse=True)]
#     return sorted_strings[:limit]
    

def search_patients(search_name, flag, limit=30):
    with open('./static/config/table_names.json', encoding='utf-8') as conf:
        tables_name = json.load(conf)
    with open('./static/config/excluded_features.json', encoding='utf-8') as conf:
        excluded_features = json.load(conf)
    with open('./static/config/shown_headers.json', encoding='utf-8') as conf:
        shown_headers = json.load(conf)
    tables_result = []
    try:
        connection = pymysql.connect(
            host=config['mysql_url'],  # 修改为你的数据库主机
            user=config['mysql_user'],  # 修改为你的数据库用户名
            password=config['mysql_password'],  # 修改为你的数据库密码
            database=config['database'],  # 修改为你的数据库名
            charset='utf8mb4'
        )
        query_subject_ids = f"""
            SELECT subject_id, `首页出院主要诊断`, `首页出院其他诊断`
            FROM {config['database']}.电子病历
            WHERE (`首页出院主要诊断` LIKE %s AND `首页出院主要诊断` LIKE %s)
            """ if flag == '0' else f"""
            SELECT subject_id, `首页出院主要诊断`, `首页出院其他诊断`
            FROM {config['database']}.电子病历
            WHERE (`首页出院主要诊断` LIKE %s OR `首页出院其他诊断` LIKE %s)
            """ 
        ids_names_dict = {}
        main_others_dict = {}
        medical_words = get_medical_words_through_models(config["ip"], search_name, config["model"])
        medical_words = medical_words if medical_words else search_name
        for sub_search_name in get_sub_words(medical_words):
            with connection.cursor() as cursor:
                cursor.execute(query_subject_ids, ((f'%{sub_search_name}%'), (f'%{sub_search_name}%')))
                ids_names = cursor.fetchall()
                ids_names_dict.update({id_name[0]:id_name[1] for id_name in ids_names}) if flag == '0' else main_others_dict.update({id_name[0]:(id_name[1],id_name[2]) for id_name in ids_names})          
        if not ids_names_dict and not main_others_dict:
            return {"error": str(e)}
        
        # results = process.extract(search_name, list(ids_names_dict.keys()), scorer=fuzz.ratio, limit=limit)
        results = sorted([(id, choice, fuzz.ratio(medical_words, choice)) for id, choice in list(ids_names_dict.items())], key=lambda x: x[2], reverse=True)[:limit] if  flag == '0' else sorted([(id, choice, fuzz.ratio(medical_words, choice[0]) + 0.5 * fuzz.ratio(medical_words, choice[1])) for id, choice in list(main_others_dict.items())], key=lambda x: x[2], reverse=True)[:limit]
        # results = get_most_similar_strings(search_name, list(ids_names_dict.keys()), limit)
        subject_id_list = [result[0] for result in results]
        def query_table(table_name, ids):
            placeholders = ', '.join(['%s'] * len(ids))
            query = f"SELECT * FROM {config['database']}.{table_name} WHERE subject_id IN ({placeholders})"
            with connection.cursor() as cursor:
                cursor.execute(query, ids)
                result = cursor.fetchall()
                columns = [desc[0] for desc in cursor.description]
            result_df = pd.DataFrame(result, columns=columns)
            # 此句是为了重新查询到的每一个病人的所有数据，可以再一次按照之前按文字相关性排好的顺序排列
            result_df['subject_id'] = pd.Categorical(result_df['subject_id'], categories=subject_id_list, ordered=True)
            result_df = result_df.sort_values('subject_id').astype(str)
            return result_df

        for table_name in tables_name:
            tables_result.append(query_table(table_name, subject_id_list))
        combined_tables_result = pd.concat(tables_result, axis=1)
        combined_tables_result = combined_tables_result.drop(columns=excluded_features)
        combined_tables_result = combined_tables_result.loc[:, ~combined_tables_result.columns.duplicated()]
        headers = list(combined_tables_result.columns)
        combined_tables_result = combined_tables_result.fillna('无')
        combined_tables_result = combined_tables_result.replace({"NaN": "无", "None": "无"})
        return {"order_headers": headers,
                "shown_headers": shown_headers,
                "patients": [OrderedDict(row) for row in combined_tables_result.to_dict(orient='records')]}
    except Exception as e:
        return {"error": str(e)}


@app.route('/patient', methods=['GET'])
def patient():
    search_name = request.args.get('search_name')
    flag = request.args.get('flag')  # 用来记录直接查找搜索框病历（1），还是节点病历（0）
    data = search_patients(search_name, flag, limit=30)
    if "error" in data:
        return jsonify(data), 500
    return jsonify(data)


@app.route('/update_node', methods=['POST'])
def update_node():
    data = request.json
    node_id = data['identity']  # 当前节点的 ID
    new_name = data['name']
    new_description = data['description']
    add_node_types = data['add_node_types']  # 新增的类别
    remove_node_types = data['remove_node_types']  # 删除的类别
    add_relationships = data['add_relationships']  # 新增的关系信息
    remove_relationships = data['remove_relationships']  # 删除的关系信息

    def update_node_transaction(tx, node_id, new_name, new_description, add_node_types, remove_node_types, add_relationships, remove_relationships):
        # 更新节点名称和简介
        update_query = """
        MATCH (n)
        WHERE ID(n) = $node_id
        SET n.name = $new_name, n.description = $new_description
        RETURN n
        """
        tx.run(update_query, node_id=node_id, new_name=new_name, new_description=new_description)

        # 处理删除类别关系，并更新types字段
        for remove_type in remove_node_types:
            remove_relationship_query = """
            MATCH (n)<-[r:包含]-(t {name: $type_name})
            WHERE ID(n) = $node_id
            DELETE r
            """
            tx.run(remove_relationship_query, node_id=node_id, type_name=remove_type)

            # 更新节点的 types 字段，删除被移除的类别
            update_types_query = """
            MATCH (n)
            WHERE ID(n) = $node_id
            SET n.types = [type IN n.types WHERE type <> $type_name]
            """
            tx.run(update_types_query, node_id=node_id, type_name=remove_type)

        # 处理新增类别关系，并更新types字段
        for add_type in add_node_types:
            # 查找目标类别节点，并创建关系
            create_relationship_query = """
            MATCH (t {name: $type_name}), (n)
            WHERE ID(n) = $node_id
            MERGE (t)-[:包含]->(n)
            """
            tx.run(create_relationship_query, node_id=node_id, type_name=add_type)

            # 更新节点的 types 字段，添加新类别
            update_types_query = """
            MATCH (n)
            WHERE ID(n) = $node_id
            SET n.types = coalesce(n.types, []) + $type_name
            """
            tx.run(update_types_query, node_id=node_id, type_name=add_type)

        # 处理关系信息
        # 处理新增关系
        for relationship in add_relationships:
            direction = relationship['direction']
            relationship_type = relationship['relationship']
            target_node_name = relationship['node']

            if direction == 'from_new_to_target':
                print("from_new_to_target:", direction)
                rel_query = """
                MATCH (n), (m {name: $target_node_name})
                WHERE ID(n) = $node_id
                MERGE (n)-[r:%s]->(m)
                """ % relationship_type
            elif direction == 'from_target_to_new':
                print("from_target_to_new:", direction)
                rel_query = """
                MATCH (n), (m {name: $target_node_name})
                WHERE ID(n) = $node_id
                MERGE (m)-[r:%s]->(n)
                """ % relationship_type
            else:
                print("errrrror:", direction)

            tx.run(rel_query, node_id=node_id, target_node_name=target_node_name)

        # 处理删除关系
        for relationship in remove_relationships:
            direction = relationship['direction']
            relationship_type = relationship['relationship']
            target_node_name = relationship['node']

            if direction == 'from_new_to_target':
                rel_query = """
                MATCH (n)-[r:%s]->(m {name: $target_node_name})
                WHERE ID(n) = $node_id
                DELETE r
                """ % relationship_type
            elif direction == 'from_target_to_new':
                rel_query = """
                MATCH (m {name: $target_node_name})-[r:%s]->(n)
                WHERE ID(n) = $node_id
                DELETE r
                """ % relationship_type
            else:
                print("errrrror:", direction)

            tx.run(rel_query, node_id=node_id, target_node_name=target_node_name)

    # 使用事务进行更新
    with driver.session() as session:
        session.write_transaction(update_node_transaction, node_id, new_name, new_description, add_node_types, remove_node_types, add_relationships, remove_relationships)

    return jsonify({"status": "success"})


# Neo4j 数据库操作函数
def add_node_with_relationship(label, name, description=None, target_name=None, relationship_type=None, relationship_direction=None):
    """
    向Neo4j数据库中添加一个新的节点，并与另一个节点建立关系。
    """
    try:
        with driver.session() as session:
            types = []
            types.append(label)
            # 检查相同类别下是否已存在相同名称的节点
            check_query = f"MATCH (n:{label}) WHERE n.name = $name RETURN n"
            result = session.run(check_query, name=name)
            existing_node = result.single()

            if existing_node is not None:
                return {"error": f"节点 {name} 在类别 {label} 下已存在"}

            # 创建新节点
            create_query = f"CREATE (n:{label} {{name: $name, description: $description, types: $types}}) RETURN n"
            result = session.run(create_query, name=name, description=description, types=types)
            new_node_record = result.single()
            if new_node_record is None:
                return {"error": "创建新节点失败"}

            new_node_id = new_node_record["n"].id

            if target_name and relationship_type:
                # 构造目标节点查询字符串
                find_query = f"MATCH (m) WHERE m.name = $target_name RETURN m"
                result = session.run(find_query, target_name=target_name)
                target_node_record = result.single()

                if target_node_record is None:
                    return {"error": "目标节点未找到"}

                target_node = target_node_record["m"]
                target_node_id = target_node.id

                # 创建关系
                if relationship_direction == "from_new_to_target":
                    rel_query = f"MATCH (n) WHERE id(n) = {new_node_id} MATCH (m) WHERE id(m) = {target_node_id} CREATE (n)-[r:{relationship_type}]->(m) RETURN r"
                elif relationship_direction == "from_target_to_new":
                    rel_query = f"MATCH (n) WHERE id(n) = {new_node_id} MATCH (m) WHERE id(m) = {target_node_id} CREATE (m)-[r:{relationship_type}]->(n) RETURN r"
                else:
                    return {"error": "无效的关系方向"}

                session.run(rel_query)

                # 返回新节点信息
                node_properties = {key: value for key, value in new_node_record["n"].items()}
                return {
                    "id": new_node_id,
                    "label": label,
                    "types": types,
                    "properties": node_properties,
                    "relationship_type": relationship_type,
                    "target_node_id": target_node_id
                }
            else:
                # 返回新节点信息
                node_properties = {key: value for key, value in new_node_record["n"].items()}
                return {
                    "id": new_node_id,
                    "label": label,
                    "properties": node_properties
                }

    except Exception as e:
        return {"error": str(e)}

# 搜索节点的路由
@app.route('/search_nodes', methods=['GET'])
def search_nodes():
    """
    根据节点名称搜索节点。
    """
    name = request.args.get('name')
    if not name:
        return jsonify({"error": "缺少名称参数"}), 400

    try:
        with driver.session() as session:
            query = f"MATCH (n) WHERE n.name CONTAINS $name RETURN n LIMIT 30"
            result = session.run(query, name=name)
            nodes = [record["n"].as_dict() for record in result]
            return jsonify(nodes)

    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/search_relationship_types', methods=['GET'])
def search_relationship_types():
    keyword = request.args.get('keyword')
    if not keyword:
        return jsonify({"error": "Keyword is required"}), 400

    # 查询的逻辑
    query = "MATCH ()-[r]->() RETURN DISTINCT type(r) AS relationshipType LIMIT 30"
    types = []
    try:
        with driver.session() as session:
            result = session.run(query)
            for record in result:
                if keyword in record['relationshipType']:
                    types.append(record['relationshipType'])
    except Exception as e:
        return jsonify({"error": str(e)}), 500

    return jsonify(types), 200

# 添加节点并建立关系的路由
@app.route('/add_node_with_relationship', methods=['POST'])
def add_node_with_relationship_api():
    data = request.form  # 使用 form 数据
    label = data.get('label')
    name = data.get('name')
    description = data.get('description')
    target_name = data.get('target_name')
    relationship_type = data.get('relationship_type')
    relationship_direction = data.get('relationship_direction')

    if not label or not name:
        return jsonify({"error": "缺少必要的参数（标签或名称）"}), 400

    # 调用 Neo4j 数据库操作函数
    result = add_node_with_relationship(label, name, description, target_name, relationship_type, relationship_direction)

    if 'error' in result:
        return jsonify(result), 400

    return jsonify(result), 201


 # 删除节点
# 删除指定节点和相关关系
def delete_node_and_relationships(node_id):
    query = """
    MATCH (n) WHERE ID(n) = $node_id
    DETACH DELETE n
    """
    with driver.session() as session:
        session.run(query, node_id=int(node_id))

@app.route('/get_related_nodes', methods=['POST'])
def get_related_nodes():
    data = request.json
    node_id = data.get('nodeId')

    if not node_id:
        return jsonify({"error": "Node ID is required"}), 400

    query = """
    MATCH (n)-[r]-(m)
    WHERE id(n) = $node_id
    RETURN m.name AS relatedNode, type(r) AS relationship,
           CASE WHEN id(startNode(r)) = $node_id THEN 'from_new_to_target' ELSE 'from_target_to_new' END AS direction
    """
    try:
        with driver.session() as session:
            results = session.run(query, node_id=int(node_id))
            related_nodes = [
                {
                    "node": record["relatedNode"],
                    "relationship": record["relationship"],
                    "direction": record["direction"]
                }
                for record in results
            ]
        return jsonify({"related_nodes": related_nodes})
    except Exception as e:
        return jsonify({"error": f"Failed to fetch related nodes: {str(e)}"}), 500
    
# 定义Flask路由处理前端请求
@app.route('/delete_node', methods=['POST'])
def delete_node():
    data = request.json
    node_id = data.get('nodeId')
    print("此时删除的节点ID为"+node_id)
    if not node_id:
        return jsonify({'message': 'Node ID is required.'}), 400

    try:
        # 删除节点及其相关关系
        delete_node_and_relationships(node_id)
        return jsonify({'message': '节点及其相关关系删除成功！'})
    except Exception as e:
        return jsonify({'message': f'Error: {str(e)}'}), 500


def format_relationships(relationships, nodes):
    # 构建一个节点ID到节点信息的映射，方便快速访问节点信息
    node_map = {node['id']: node for node in nodes}

    # 存储格式化后的关系字符串
    formatted_relationships = []

    # 遍历每个关系
    for rel in relationships:
        # 找到关系对应的节点a和节点b的信息
        node_a = node_map.get(rel['start'])
        node_b = node_map.get(rel['end'])

        # 检查节点a和节点b是否存在
        if node_a and node_b:
            # 格式化为“（实体a名称，实体a和b的相关关系，实体b的名称）”形式
            formatted_str = f"({node_a['properties']['name']}, {rel['type']}, {node_b['properties']['name']})"
            formatted_relationships.append(formatted_str)

    return formatted_relationships

@app.route('/extend')
def extend():
    # 返回 AI助手页面的 HTML
    return render_template('extend.html')

@app.route('/transformation', methods=['GET', 'POST'])
def ask_question():
    if request.method == 'POST':
        question = request.form['question']
        graph_data = get_graph_data(question)
        # 格式化关系数据
        formatted_relationships = format_relationships(graph_data['relationships'], graph_data['nodes'])

        # 将格式化后的关系数据每行一条拼接在一起
        formatted_str = "\n".join(formatted_relationships)
        extended_question = f"{question} 从知识图谱中检索到的三元组如下：{formatted_str}"

        print(extended_question)
        return jsonify({
            'question': question,
            'extended_question': extended_question
        })
    return jsonify({'message': 'Please submit a question.'})

def ask_question(question):
    graph_data = get_graph_data(question)
    # 格式化关系数据
    formatted_relationships = format_relationships(graph_data['relationships'], graph_data['nodes'])

    # 将格式化后的关系数据每行一条拼接在一起
    formatted_str = "\n".join(formatted_relationships)
    extended_question = f"{question} 从知识图谱中检索到的三元组如下：{formatted_str}"
    print(extended_question)
    return extended_question
    # if formatted_str:
    #     extended_question = f"{question} 从知识图谱中检索到的三元组如下：{formatted_str}"
    #     print(extended_question)
    #     return extended_question
    # else:
    #     extended_question = question
    #
    # print(extended_question)
    # return extended_question

@app.route('/AI')
def ai_helper():
    # 返回 AI助手页面的 HTML
    return render_template('AI.html')


# 辅助函数，用于获取最近的 N 个用户的问题
def get_recent_user_questions(messages, n):
    user_questions = []  # 存储用户的问题
    for message in reversed(messages):  # 从后往前遍历，找到n个用户的问题
        if message['role'] == 'user':  # 当遇到用户的问题时
            user_questions.insert(0, message['questions'])  # 将用户的问题添加到列表中
            if len(user_questions) == n:  # 如果已经找到了n个问题
                break
    return user_questions

def generate_suggestions_through_models(input_content):
    url = config['base_url']
    full_content = f"从医学的角度出发，请用中文回答，我现在会给你一些用户之前问过的问题，请根据这些问题，预测接下来两个问题，问题之间使用、分开，请只返回这两个问题，不用给序号、不要解释、。之前的问题如下：'{input_content}'"
    data = {
        "model": config['model'],
        "messages": [{
            "role": "system",
            "content": full_content
        }],
        "stream": True
    }

    headers = config['headers']

    # url = f"{config['base_url']}/chat"
    #
    # headers = config['headers']
    #
    # full_content = f"从医学的角度出发，请用中文回答，我现在会给你一些用户之前问过的问题，请根据这些问题，预测接下来两个问题，问题之间使用、分开，请只返回这两个问题，不用给序号、不要解释、。之前的问题如下：'{input_content}'"
    #
    # data = {
    #     "model": config['model'],
    #     "messages": [{
    #         "role": "system",
    #         "content": full_content
    #     }]
    # }

    # host = "172.18.5.119"
    # port = "11434"
    # url = f"http://{ip}:{port}/api/chat"
    # model = "qwen2.5:3b"
    # model = "wangshenzhi/llama3-8b-chinese-chat-ollama-q4"
    # headers = {"Content-Type": "application/json"}
    # full_content = f"从医学的角度出发，请用中文回答，我现在会给你一些用户之前问过的问题，请根据这些问题，请预测接下来可能被问到的两个问题，只返回这两个问题本身，不要解释，不要序号 问题如下：‘{input_content}‘"
    # full_content = f"基于医学知识，请用中文预测接下来可能被问到的两个问题，请使用‘、’分隔问题，问题前面不需要序号，并且不要提供解释。以下是用户之前提出的问题：‘{input_content}’"

    # data = {
    #     "model": model,
    #     "options": {
    #         "temperature": 0.
    #     },
    #     "stream": False,
    #     "messages": [{
    #         "role": "system",
    #         "content": full_content
    #     }]
    # }
    try:
        response = requests.post(url, json=data, headers=headers, timeout=20)
        # 检查响应状态码
        if response.status_code == 200:
            # 解析JSON
            res = response.json()
            # 直接提取content部分
            content = res['message']['content']
            # 假设预测的问题使用'、'分隔
            questions = content.split('\n')
            # 返回前两个问题
            two_questions = questions[:2]
            # cleaned_questions = [re.sub(r'^\W+|\W+$', '', q) for q in two_questions]
            cleaned_questions = [re.sub(r'^\d+\.\s*|\W+$', '', q) + '?' for q in  two_questions]
            return cleaned_questions  # 返回模型预测的两个问题
        else:
            return []  # 如果状态码不是200，返回空列表
    except requests.exceptions.Timeout:
        return []  # 如果请求超时，返回空列表
    except ValueError as e:
        # 如果解析JSON失败，打印错误信息并返回空列表
        print("解析JSON失败：", e)
        return []
    except requests.exceptions.RequestException as e:
        # 如果发生其他请求异常，打印错误信息并返回空列表
        print("请求异常：", e)
        return []

@app.route('/api/chat', methods=['POST'])
def chat():
    global conversation
    message = request.json.get('message')
    # 进入知识图谱搜索三元组,扩展问题
    extended_message = ask_question(message)
    conversation.append({"role": "user", "content": extended_message, "questions": message})


    stream = ollama.chat(
    model='qwen2.5:3b',
    messages=conversation,
    stream=True
    )

    result = ""

    for chunk in stream:
        print(chunk['message']['content'], end='', flush=True)
        result += chunk['message']['content']

    conversation.append({"role": "assistant", "content": result, "questions": result})

    # resp_message = response.choices[0].delta

    # ---  添加问题联想功能 ---
    # 获取最近的1个完整对话
    recent_questions = get_recent_user_questions(conversation, 1)

    # # 获取实际对话个数
    # if(len(recent_dialogues) > 0) :
    #     num_dialogues = min(len(recent_dialogues), 3)
    # recent_dialogues = recent_dialogues[:num_dialogues]

    # 构建 prompt
    prompt = "、".join(q + "?" for q in recent_questions)

    print("构建的大模型输入prompt：", prompt)

    # 调用服务器模型进行问题联想
    suggestions = generate_suggestions_through_models(prompt)
    print("联想出的问题:", suggestions)

    resp_message = result
    return jsonify({"response": resp_message, "suggestions": suggestions})




# @app.route('/api/chat', methods=['POST'])
# def chat():
#     global conversation
#     message = request.json.get('message')
#     # 进入知识图谱搜索三元组,扩展问题
#     extended_message = ask_question(message)
#     conversation.append({"role": "user", "content": extended_message})



#     response = client.chat.completions.create(
#         model="glm-4-flash",  # 填写需要调用的模型编码
#         messages=conversation,
#         stream=True,
#     )
#     result = ""
#     for chunk in response:
#         print(chunk.choices[0].delta)
#         result += chunk.choices[0].delta.content
#     # print(response.choices[0].message)

#     conversation.append({"role": "assistant", "content": result})

#     # resp_message = response.choices[0].delta

#     resp_message = result
#     return jsonify({"response": resp_message})


@app.route('/chat/sse', methods=['POST'])
def ssechat():
    messages =  request.get_json()['messages']
    all_messages = ""
    for message in messages:
        if message['role'] == 'user':
            all_messages += " "
            all_messages += message['content']
    extended_message = ask_question(all_messages)
    # messages.append({"role": "assistant", "content": ""})
    messages.append({"role": "user", "content": extended_message})

    url = config['base_url']

    data = {
        "model": config['model'],
        "messages": messages,
        "stream": True
    }

    headers = config['headers']

    response = requests.post(url, headers=headers, json=data, stream=True)

    def generate():
        for chunk in response:
            # print(chunk)
            yield chunk
            # yield f"data: {chunk}\n\n"
            # serializable_data = serialize_chat_completion_chunk(chunk)
            # yield f"data: {json.dumps(serializable_data, ensure_ascii=False)}\n\n"
        # yield "data: [DONE]\n\n"
    return Response(generate(), mimetype='text/event-stream')

# def serialize_chat_completion_chunk(chunk):
#     # 假设 chunk 是一个包含可序列化字段的对象
#     return {
#         'id': chunk.id,
#         'created': chunk.created,
#         'model': chunk.model,
#         'choices': [
#             {
#                 'index': choice.index,
#                 'delta': {
#                     'role': choice.delta.role,
#                     'content': choice.delta.content,
#                 },
#             }
#             for choice in chunk.choices
#         ],
#     }


# 问题联想接口部分，里面有内设的dialogue进行测试，可以把dialogue重制为传入的参数进行实现。
@app.route('/question/similarity', methods=['POST'])
def get_answer():

    # client = ZhipuAI(api_key="bbcb6772972043f3b8e4c28fcdeeceeb.gvqbdyK7lckWNdOf")
    # 传入并修改dialogue部分
    dialogue = """
    用户: 麻醉有哪些注意事项  
    ai助手: 麻醉的注意事项包括但不限于以下几点：
    ### 麻醉前的注意事项
    - **健康状况告知**：向医生详细告知自己的健康状况，包括既往病史、过敏史、正在服用的药物、有无吸烟、饮酒等习惯。
    - **禁食禁饮**：通常需要在麻醉前6-8小时内禁食，2-4小时内禁饮，以减少麻醉过程中呕吐和误吸的风险。
    - **个人物品**：取下所有首饰、假牙、隐形眼镜等，以免在手术过程中造成损伤或干扰手术。
    - **心理准备**：了解麻醉的必要性和安全性，消除紧张和恐惧心理，保持良好的心态。
    - **签署知情同意书**：了解麻醉的风险和可能的并发症，并签署知情同意书。

    ### 麻醉中的注意事项
    - **配合医生**：在麻醉过程中，尽量保持平静，配合医生的指示，如有不适或异常感觉，及时告知医生。
    - **监测生命体征**：医生会持续监测患者的生命体征，如心率、血压、血氧饱和度等，以确保麻醉的安全。
    - **呼吸管理**：对于全身麻醉的患者，医生会进行呼吸管理，确保患者的呼吸通畅和氧气供应。

    """

    try:
        url = config['base_url']

        headers = config['headers']

        data = {
            "model": config['model'],
            "messages": [
                {
                    "role": "user",
                    "content": f"根据以下对话生成与之相关的3个问题：\n\n对话内容：{dialogue}\n\n生成的问题（给出问题前面加个数字序号）："
                }   # 提问的内容
            ]
        }

        response = requests.post(url=url, headers=headers, json=data)
        # response = client.chat.completions.create(
        #     model="glm-4",  # 使用的模型，根据实际情况选择合适的模型
        #     messages=[
        #         {"role": "user",
        #          "content": f"根据以下对话生成与之相关的3个问题：\n\n对话内容：{dialogue}\n\n生成的问题（给出问题前面加个数字序号）："
        #          }   # 提问的内容
        #     ],
        # )
        # return response

        response = json.loads(response.text)
        answer = response['choices'][0]['message']['content'] if response['choices'] else ""
        print(answer)
        # 假设返回的内容是以'、'分隔的多个问题
        questions = answer.split('\n')

        # 去掉每个问题前面的序号（例如 '1. '），使用正则表达式
        cleaned_questions = [re.sub(r'^\d+\.\s*', '', question).strip() for question in questions]
        print(cleaned_questions)
        # 返回去掉序号后的问题列表
        return cleaned_questions[:3]  # 返回前三个问题
    except APIStatusError as err:
        print(f"请求失败，错误信息：{err}")
        return []
    except Exception as e:
        print(f"请求过程中出现其他错误: {e}")
        return []


if __name__ == '__main__':
    app.run(host='0.0.0.0',debug=True, port=8000)
